/******************************************************************************
 * Copyright 2022 The Untouch Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/
#include "modules/calibration/lidar_imu_calibrator/include/odometry_frame.h"

#include <glk/primitives/primitives.hpp>
#include <pcl/filters/voxel_grid.h>
#include "pcl/io/pcd_io.h"

namespace lidar_imu_calibrator {

OdometryFrame::OdometryFrame(const std::string& raw_cloud_path,
                             const Eigen::Isometry3d& pose, double timestamp)
    : point_cloud_buffer_need_reset_(false),
      timestamp_(timestamp),
      pose_(pose),
      pose_imu_(pose),
      raw_cloud_path_(raw_cloud_path),
      downsample_resolution_(0.2f),
      point_cloud_(nullptr) {}

double OdometryFrame::GetFrameTimestamp() { return timestamp_; }

const Eigen::Isometry3d& OdometryFrame::GetFramePose() { return pose_; }

const Eigen::Isometry3d& OdometryFrame::GetImuFramePose() { return pose_imu_; }


void OdometryFrame::SetFrameImuPose(const Eigen::Isometry3d& pose) {
  pose_imu_ = pose;
}

void OdometryFrame::UpdateFramePose(const Eigen::Isometry3d& extrinsics) {
  pose_ = pose_imu_ * extrinsics;
}

void OdometryFrame::LoadPointCloud(
    const std::vector<double>& timestamps,
    std::vector<Eigen::Isometry3d, Eigen::aligned_allocator<Eigen::Isometry3d>>&
        poses,
    Eigen::Isometry3d& extrinsic) {
  pcl::PointCloud<velodyne::PointXYZIT>::Ptr input_point_cloud(
      new pcl::PointCloud<velodyne::PointXYZIT>());
  Eigen::Isometry3d pose_min_time;
  Eigen::Isometry3d pose_max_time;
  double timestamp_min;
  double timestamp_max;

  pcl::io::loadPCDFile(raw_cloud_path_, *input_point_cloud);
  pcl::PointCloud<pcl::PointXYZI>::Ptr compensate_cloud(
      new pcl::PointCloud<pcl::PointXYZI>());
  compensate_cloud->width = input_point_cloud->width;
  compensate_cloud->height = input_point_cloud->height;
  compensate_cloud->is_dense = input_point_cloud->is_dense;
  compensate_cloud->points.resize(input_point_cloud->size());

  GetTimestampInterval(input_point_cloud, timestamp_min, timestamp_max);
  ComputeInterpolationPose(timestamp_min, timestamps, poses, extrinsic,
                           pose_min_time);
  ComputeInterpolationPose(timestamp_max, timestamps, poses, extrinsic,
                           pose_max_time);
  MotionCompensation(input_point_cloud, timestamp_min, timestamp_max,
                     pose_min_time, pose_max_time, compensate_cloud);

  pcl::PointCloud<pcl::PointXYZI>::Ptr downsampled(
      new pcl::PointCloud<pcl::PointXYZI>());

  pcl::VoxelGrid<pcl::PointXYZI> voxel_grid;
  voxel_grid.setLeafSize(downsample_resolution_, downsample_resolution_,
                         downsample_resolution_);
  voxel_grid.setInputCloud(compensate_cloud);
  voxel_grid.filter(*downsampled);

  point_cloud_ = downsampled;
  point_cloud_buffer_need_reset_ = true;
}

void OdometryFrame::GetPointCloud() {
  if (point_cloud_buffer_ == nullptr || point_cloud_buffer_need_reset_) {
    point_cloud_buffer_.reset(new glk::PointCloudBuffer(point_cloud_));
    point_cloud_buffer_need_reset_ = false;
  }
}

void OdometryFrame::GetPointCloud(pcl::PointCloud<pcl::PointXYZI>::Ptr &point_cloud){
  point_cloud->clear();
  if (point_cloud_->empty()) return;
  pcl::copyPointCloud(*point_cloud_, *point_cloud);
}

void OdometryFrame::DrawPointCloud(glk::GLSLShader& shader,
                                   float downsample_resolution) {
  GetPointCloud();

  shader.set_uniform("color_mode", 0);
  shader.set_uniform("model_matrix", pose_.cast<float>().matrix());
  point_cloud_buffer_->draw(shader);

  shader.set_uniform("color_mode", 1);
  shader.set_uniform("material_color", Eigen::Vector4f(1.0f, 0.0f, 0.0f, 1.0f));
  shader.set_uniform("apply_keyframe_scale", true);
  auto& sphere =
      glk::Primitives::instance()->primitive(glk::Primitives::SPHERE);
  sphere.draw(shader);
  shader.set_uniform("apply_keyframe_scale", false);
}

void OdometryFrame::GetTimestampInterval(
    pcl::PointCloud<velodyne::PointXYZIT>::Ptr& input_point_cloud,
    double& timestamp_min, double& timestamp_max) {
  timestamp_max = 0;
  timestamp_min = std::numeric_limits<uint64_t>::max();

  for (const auto& point : input_point_cloud->points) {
    if (point.timestamp < timestamp_min) {
      timestamp_min = point.timestamp;
    }

    if (point.timestamp > timestamp_max) {
      timestamp_max = point.timestamp;
    }
  }
}

size_t OdometryFrame::GetIndex(const std::vector<double>& timestamps,
                               const double ref_timestamp) {
  size_t index =
      std::lower_bound(timestamps.begin(), timestamps.end(), ref_timestamp) -
      timestamps.begin();
  index = std::max(1lu, std::min(index, timestamps.size() - 1));

  return index;
}

void OdometryFrame::ComputeInterpolationPose(
    double timestamp, const std::vector<double>& timestamps,
    std::vector<Eigen::Isometry3d, Eigen::aligned_allocator<Eigen::Isometry3d>>&
        in_poses,
    Eigen::Isometry3d& extrinsic, Eigen::Isometry3d& out_pose) {
  size_t index = GetIndex(timestamps, timestamp);
  double cur_timestamp = timestamps[index];
  double pre_timestamp = timestamps[index - 1];
  double t = (cur_timestamp - timestamp) / (cur_timestamp - pre_timestamp);
  if (t < 0.0 || t > 1.0) {
    out_pose = in_poses[index] * extrinsic;
  }

  Eigen::Isometry3d pre_pose = in_poses[index - 1];
  Eigen::Isometry3d cur_pose = in_poses[index];
  Eigen::Quaterniond pre_quatd(pre_pose.linear());
  Eigen::Translation3d pre_transd(pre_pose.translation());
  Eigen::Quaterniond cur_quatd(cur_pose.linear());
  Eigen::Translation3d cur_transd(cur_pose.translation());

  Eigen::Quaterniond res_quatd = pre_quatd.slerp(1 - t, cur_quatd);

  Eigen::Translation3d re_transd;
  re_transd.x() = pre_transd.x() * t + cur_transd.x() * (1 - t);
  re_transd.y() = pre_transd.y() * t + cur_transd.y() * (1 - t);
  re_transd.z() = pre_transd.z() * t + cur_transd.z() * (1 - t);

  out_pose = re_transd * res_quatd * extrinsic;
}

void OdometryFrame::MotionCompensation(
    pcl::PointCloud<velodyne::PointXYZIT>::Ptr& input_point_cloud,
    double timestamp_min, double timestamp_max,
    const Eigen::Isometry3d& pose_min_time,
    const Eigen::Isometry3d& pose_max_time,
    pcl::PointCloud<pcl::PointXYZI>::Ptr& output_point_cloud) {
  Eigen::Vector3d translation =
      pose_min_time.translation() - pose_max_time.translation();
  Eigen::Quaterniond q_max(pose_max_time.linear());
  Eigen::Quaterniond q_min(pose_min_time.linear());
  Eigen::Quaterniond q1(q_max.conjugate() * q_min);
  Eigen::Quaterniond q0(Eigen::Quaterniond::Identity());
  q1.normalize();
  translation = q_max.conjugate() * translation;

  double d = q0.dot(q1);
  double abs_d = std::abs(d);
  double f = 1.0 / static_cast<double>(timestamp_max - timestamp_min);

  if (abs_d < 1.0 - 1.0e-8) {
    double theta = std::acos(abs_d);
    double sin_theta = std::sin(theta);
    double c1_sign = (d > 0) ? 1 : -1;
    for (size_t i = 0; i < output_point_cloud->points.size(); ++i) {
      float x_scalar = input_point_cloud->points[i].x;
      if (std::isnan(x_scalar)) {
        output_point_cloud->points[i].x = input_point_cloud->points[i].x;
        output_point_cloud->points[i].y = input_point_cloud->points[i].y;
        output_point_cloud->points[i].z = input_point_cloud->points[i].z;
        output_point_cloud->points[i].intensity =
            input_point_cloud->points[i].intensity;
        continue;
      }

      float y_scalar = input_point_cloud->points[i].y;
      float z_scalar = input_point_cloud->points[i].z;
      Eigen::Vector3d p(x_scalar, y_scalar, z_scalar);
      double t = (timestamp_max - input_point_cloud->points[i].timestamp) * f;

      Eigen::Translation3d ti(t * translation);

      double c0 = std::sin((1 - t) * theta) / sin_theta;
      double c1 = std::sin(t * theta) / sin_theta * c1_sign;
      Eigen::Quaterniond qi(c0 * q0.coeffs() + c1 * q1.coeffs());
      Eigen::Affine3d trans = ti * qi;
      p = trans * p;

      output_point_cloud->points[i].x = static_cast<float>(p.x());
      output_point_cloud->points[i].y = static_cast<float>(p.y());
      output_point_cloud->points[i].z = static_cast<float>(p.z());
      output_point_cloud->points[i].intensity =
          input_point_cloud->points[i].intensity;
    }
  } else {
    for (size_t i = 0; i < output_point_cloud->points.size(); ++i) {
      float x_scalar = input_point_cloud->points[i].x;
      if (std::isnan(x_scalar)) {
        output_point_cloud->points[i].x = input_point_cloud->points[i].x;
        output_point_cloud->points[i].y = input_point_cloud->points[i].y;
        output_point_cloud->points[i].z = input_point_cloud->points[i].z;
        output_point_cloud->points[i].intensity =
            input_point_cloud->points[i].intensity;
        continue;
      }
      float y_scalar = input_point_cloud->points[i].y;
      float z_scalar = input_point_cloud->points[i].z;
      Eigen::Vector3d p(x_scalar, y_scalar, z_scalar);

      double t = (timestamp_max - input_point_cloud->points[i].timestamp) * f;
      Eigen::Translation3d ti(t * translation);

      p = ti * p;

      output_point_cloud->points[i].x = static_cast<float>(p.x());
      output_point_cloud->points[i].y = static_cast<float>(p.y());
      output_point_cloud->points[i].z = static_cast<float>(p.z());
      output_point_cloud->points[i].intensity =
          input_point_cloud->points[i].intensity;
    }
  }
}
}  // namespace lidar_imu_calibrator
